diff --git a/arch/arm/include/asm/arch-rockchip/boot_mode.h b/arch/arm/include/asm/arch-rockchip/boot_mode.h
index bd65f60..163b2e7 100644
--- a/arch/arm/include/asm/arch-rockchip/boot_mode.h
+++ b/arch/arm/include/asm/arch-rockchip/boot_mode.h
@@ -16,4 +16,6 @@
 /* enter usb mass storage mode */
 #define BOOT_UMS		(REBOOT_FLAG + 12)
 
+int setup_boot_mode(void);
+
 #endif
diff --git a/arch/arm/mach-rockchip/Kconfig b/arch/arm/mach-rockchip/Kconfig
index 36df484..6144057 100644
--- a/arch/arm/mach-rockchip/Kconfig
+++ b/arch/arm/mach-rockchip/Kconfig
@@ -117,6 +117,7 @@
 	select SPL_SERIAL_SUPPORT
 	select SPL_DRIVERS_MISC_SUPPORT
 	select DEBUG_UART_BOARD_INIT
+	select BOARD_LATE_INIT
 	help
 	  The Rockchip RK3399 is a ARM-based SoC with a dual-core Cortex-A72
 	  and quad-core Cortex-A53.
@@ -152,6 +153,20 @@
           SPL will return to the boot rom, which will then load the U-Boot
           binary to keep going on.
 
+config ROCKCHIP_BOOT_MODE_REG
+	hex "Rockchip boot mode flag register address"
+	default 0x200081c8 if ROCKCHIP_RK3036
+	default 0x20004040 if ROCKCHIP_RK3188
+	default 0x110005c8 if ROCKCHIP_RK322X
+	default 0xff730094 if ROCKCHIP_RK3288
+	default 0xff738200 if ROCKCHIP_RK3368
+	default 0xff320300 if ROCKCHIP_RK3399
+	default 0x10300580 if ROCKCHIP_RV1108
+	default 0
+	help
+	  The Soc will enter to different boot mode(defined in asm/arch/boot_mode.h)
+	  according to the value from this register.
+
 config ROCKCHIP_SPL_RESERVE_IRAM
 	hex "Size of IRAM reserved in SPL"
 	default 0x4000
diff --git a/arch/arm/mach-rockchip/Makefile b/arch/arm/mach-rockchip/Makefile
index c15e9bf..2127f2b 100644
--- a/arch/arm/mach-rockchip/Makefile
+++ b/arch/arm/mach-rockchip/Makefile
@@ -22,10 +22,16 @@
 obj-spl-$(CONFIG_ROCKCHIP_RK3399) += rk3399-board-spl.o spl-boot-order.o
 
 ifeq ($(CONFIG_SPL_BUILD)$(CONFIG_TPL_BUILD),)
+
+ifneq ($(CONFIG_ROCKCHIP_BOOT_MODE_REG),0)
+obj-y += boot_mode.o
+endif
+
 obj-$(CONFIG_ROCKCHIP_RK3188) += rk3188-board.o
 obj-$(CONFIG_ROCKCHIP_RK322X) += rk322x-board.o
 obj-$(CONFIG_ROCKCHIP_RK3288) += rk3288-board.o
 obj-$(CONFIG_ROCKCHIP_RK3036) += rk3036-board.o
+obj-$(CONFIG_ROCKCHIP_RK3399) += rk3399-board.o
 endif
 
 obj-$(CONFIG_$(SPL_TPL_)RAM) += sdram_common.o
diff --git a/arch/arm/mach-rockchip/boot_mode.c b/arch/arm/mach-rockchip/boot_mode.c
new file mode 100644
index 0000000..4652490
--- /dev/null
+++ b/arch/arm/mach-rockchip/boot_mode.c
@@ -0,0 +1,33 @@
+/*
+ * (C) Copyright 2016 Rockchip Electronics Co., Ltd
+ *
+ * SPDX-License-Identifier:     GPL-2.0+
+ */
+
+#include <common.h>
+#include <asm/io.h>
+#include <asm/arch/boot_mode.h>
+
+int setup_boot_mode(void)
+{
+	void *reg = (void *)CONFIG_ROCKCHIP_BOOT_MODE_REG;
+	int boot_mode = readl(reg);
+
+	debug("boot mode %x.\n", boot_mode);
+
+	/* Clear boot mode */
+	writel(BOOT_NORMAL, reg);
+
+	switch (boot_mode) {
+	case BOOT_FASTBOOT:
+		printf("enter fastboot!\n");
+		env_set("preboot", "setenv preboot; fastboot usb0");
+		break;
+	case BOOT_UMS:
+		printf("enter UMS!\n");
+		env_set("preboot", "setenv preboot; ums mmc 0");
+		break;
+	}
+
+	return 0;
+}
diff --git a/arch/arm/mach-rockchip/rk3036-board.c b/arch/arm/mach-rockchip/rk3036-board.c
index a3457f3..a5d2571 100644
--- a/arch/arm/mach-rockchip/rk3036-board.c
+++ b/arch/arm/mach-rockchip/rk3036-board.c
@@ -19,30 +19,6 @@
 
 DECLARE_GLOBAL_DATA_PTR;
 
-#define GRF_BASE	0x20008000
-
-static void setup_boot_mode(void)
-{
-	struct rk3036_grf *const grf = (void *)GRF_BASE;
-	int boot_mode = readl(&grf->os_reg[4]);
-
-	debug("boot mode %x.\n", boot_mode);
-
-	/* Clear boot mode */
-	writel(BOOT_NORMAL, &grf->os_reg[4]);
-
-	switch (boot_mode) {
-	case BOOT_FASTBOOT:
-		printf("enter fastboot!\n");
-		env_set("preboot", "setenv preboot; fastboot usb0");
-		break;
-	case BOOT_UMS:
-		printf("enter UMS!\n");
-		env_set("preboot", "setenv preboot; ums mmc 0");
-		break;
-	}
-}
-
 __weak int rk_board_late_init(void)
 {
 	return 0;
diff --git a/arch/arm/mach-rockchip/rk3188-board.c b/arch/arm/mach-rockchip/rk3188-board.c
index 96859a5..916d18f 100644
--- a/arch/arm/mach-rockchip/rk3188-board.c
+++ b/arch/arm/mach-rockchip/rk3188-board.c
@@ -24,6 +24,7 @@
 {
 	struct rk3188_grf *grf;
 
+	setup_boot_mode();
 	grf = syscon_get_first_range(ROCKCHIP_SYSCON_GRF);
 	if (IS_ERR(grf)) {
 		pr_err("grf syscon returned %ld\n", PTR_ERR(grf));
diff --git a/arch/arm/mach-rockchip/rk322x-board.c b/arch/arm/mach-rockchip/rk322x-board.c
index d443114..e71847d 100644
--- a/arch/arm/mach-rockchip/rk322x-board.c
+++ b/arch/arm/mach-rockchip/rk322x-board.c
@@ -16,30 +16,6 @@
 
 DECLARE_GLOBAL_DATA_PTR;
 
-#define GRF_BASE	0x11000000
-
-static void setup_boot_mode(void)
-{
-	struct rk322x_grf *const grf = (void *)GRF_BASE;
-	int boot_mode = readl(&grf->os_reg[0]);
-
-	debug("boot mode %x.\n", boot_mode);
-
-	/* Clear boot mode */
-	writel(BOOT_NORMAL, &grf->os_reg[0]);
-
-	switch (boot_mode) {
-	case BOOT_FASTBOOT:
-		printf("enter fastboot!\n");
-		env_set("preboot", "setenv preboot; fastboot usb0");
-		break;
-	case BOOT_UMS:
-		printf("enter UMS!\n");
-		env_set("preboot", "setenv preboot; ums mmc 0");
-		break;
-	}
-}
-
 __weak int rk_board_late_init(void)
 {
 	return 0;
diff --git a/arch/arm/mach-rockchip/rk3288-board.c b/arch/arm/mach-rockchip/rk3288-board.c
index 278bb40..1c53cca 100644
--- a/arch/arm/mach-rockchip/rk3288-board.c
+++ b/arch/arm/mach-rockchip/rk3288-board.c
@@ -23,31 +23,6 @@
 
 DECLARE_GLOBAL_DATA_PTR;
 
-#define PMU_BASE	0xff730000
-
-static void setup_boot_mode(void)
-{
-	struct rk3288_pmu *const pmu = (void *)PMU_BASE;
-	int boot_mode = readl(&pmu->sys_reg[0]);
-
-	debug("boot mode %x.\n", boot_mode);
-
-	/* Clear boot mode */
-	writel(BOOT_NORMAL, &pmu->sys_reg[0]);
-
-	switch (boot_mode) {
-	case BOOT_FASTBOOT:
-		printf("enter fastboot!\n");
-		env_set("preboot", "setenv preboot; fastboot usb0");
-		break;
-	case BOOT_UMS:
-		printf("enter UMS!\n");
-		env_set("preboot", "setenv preboot; if mmc dev 0;"
-		       "then ums mmc 0; else ums mmc 1;fi");
-		break;
-	}
-}
-
 __weak int rk_board_late_init(void)
 {
 	return 0;
diff --git a/arch/arm/mach-rockchip/rk3399-board.c b/arch/arm/mach-rockchip/rk3399-board.c
new file mode 100644
index 0000000..9293843
--- /dev/null
+++ b/arch/arm/mach-rockchip/rk3399-board.c
@@ -0,0 +1,14 @@
+/*
+ * Copyright (c) 2017 Rockchip Electronics Co., Ltd
+ *
+ * SPDX-License-Identifier:     GPL-2.0+
+ */
+
+#include <common.h>
+#include <asm/arch/boot_mode.h>
+
+int board_late_init(void)
+{
+	setup_boot_mode();
+	return 0;
+}